from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
import os
from api import UserDB, get_password_hash

# 配置数据库连接
SQLALCHEMY_DATABASE_URL = "sqlite:///" + os.path.join(os.path.dirname(os.path.abspath(__file__)), "data", "users.db")
engine = create_engine(SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# 创建数据目录
os.makedirs(os.path.join(os.path.dirname(os.path.abspath(__file__)), "data"), exist_ok=True)

# 截断密码函数，确保不超过bcrypt的72字节限制
def truncate_password(password):
    return password[:72]  # 截断到72个字符

# 初始化数据库并添加测试用户
def init_test_users():
    # 创建表
    UserDB.metadata.create_all(bind=engine)
    
    # 添加测试用户
    db = SessionLocal()
    
    # 检查admin用户是否存在
    admin_user = db.query(UserDB).filter(UserDB.username == "admin").first()
    if not admin_user:
        # 创建admin用户（管理员）
        admin_user = UserDB(
            username="admin",
            password=get_password_hash(truncate_password("admin")),  # 使用哈希密码
            is_admin=True,
            full_name="管理员",
            nickname="admin"
        )
        db.add(admin_user)
    else:
        # 更新现有admin用户的密码和权限
        admin_user.password = get_password_hash(truncate_password("admin"))
        admin_user.is_admin = True
    
    # 检查user1用户是否存在
    user1 = db.query(UserDB).filter(UserDB.username == "user1").first()
    if not user1:
        # 创建普通用户
        user1 = UserDB(
            username="user1",
            password=get_password_hash(truncate_password("password1")),
            is_admin=False,
            full_name="普通用户",
            nickname="user1"
        )
        db.add(user1)
    else:
        # 更新现有用户的密码
        user1.password = get_password_hash(truncate_password("password1"))
    
    # 提交更改
    db.commit()
    db.close()
    
    print("测试用户初始化完成！")
    print("可用账号：")
    print("- 管理员账号: admin/admin")
    print("- 普通用户账号: user1/password1")

if __name__ == "__main__":
    init_test_users()